# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com


from numpy import linspace
from copy import deepcopy

class FiniteDifferenceVariable:

    def __init__(self, name, minVal, maxVal):
        self.name = name
        self.minVal = minVal
        self.maxVal = maxVal

##class FDFunction():
##    
##    def __init__(self):
##        self.argNames = []
##        self.funct = 

class Relation:
    def __init__(self):
        self.variableNames = []
        self.numberSteps = []
        self.func = lambda x, y : 0 # x is the dist between nodes, y is the value at the node
        self.varIndices = []

    def coeffFunc(self, gridDeltas, gridVals):
        return self.func(gridDeltas, gridVals)


class BoundaryCondition:

    def __init__(self):
        self.conditionNames = []
        #self.conditionStrings = []
        self.LeftOrRight = []
        self.funct = lambda x: 0 # x is the data at the nodes


class Grid:

    def __init__(self):
        self.g = []
        self.namesOrder = []
        self.data = []
        self.ndims = 0
        self.dimSizes = []

        self.dists = []        
        
    def makeGrid(self, variables, partitions):
        self.ndims = len(variables)
        self.g = [[] for idx in range(self.ndims)]
        self.namesOrder = []
        self.data = []
        self.dimSizes = [partitions[i] + 1 for i in range(len(partitions))]
        
        if len(partitions) == self.ndims:
            tempDim = 1.0
            for idx in range(self.ndims):
                variable = variables[idx]
                partition = self.dimSizes[idx]
                dim = linspace(variable.minVal, variable.maxVal, partition)
                self.dists.append(float(variable.maxVal-variable.minVal)/partition)
                self.g[idx] = dim
                tempDim *= len(dim)
                self.namesOrder.append(variable.name)
            self.data = [0.0 for idx in range(int(tempDim))]

    def getDataIndex(self, coords):
        tempIdx = 0
        tempDimSize = 1.0
        for idx in range(len(coords)):
            coord = coords[idx]
            tempIdx += coord * tempDimSize
            tempDimSize *= self.dimSizes[idx]
        return int(tempIdx)

    def getData(self, coords):
        return self.data[self.getDataIndex(coords)]

    def setData(self, coords, val):
        self.data[self.getDataIndex(coords)] = val

    def getVariableValues(self, coords):
        returnVals = [0 for i in range(self.ndims)]
        for idx in range(len(coords)):
            returnVals[idx] = self.g[idx][coords[idx]]
        return returnVals

    def findIndexInNamesOrder(self, varName):
        for idx in range(len(self.namesOrder)):
            if self.namesOrder[idx] == varName:
                return idx

        return -1

    def getDistancesToNextNode(self, coords):    
        return self.dists
        


class FiniteDifferencesEngine:

    def __init__(self, variables, partitions, transientVariableName):
        self.grid = Grid()
        self.grid.makeGrid(variables, partitions)
        self.transientIndex = self.grid.findIndexInNamesOrder(transientVariableName)
        self.BoundaryConditionNodes = []
        self.SpecialNodes = []
        self.SpecialRelations = []
        self.mainNodeCoords = [0 for idx in range(self.grid.ndims)]

    def calcBoundaryConditionNodes(self, boundaryCondition):
        coords = [0 for idx in range(self.grid.ndims)]
        total = 1.0
        numConditions = len(boundaryCondition.conditionNames)
        variableIndices = [i for i in range(self.grid.ndims)]
        for idx in range(self.grid.ndims-1, -1, -1):
            if self.grid.namesOrder[idx] in boundaryCondition.conditionNames:
                variableIndices.remove(idx)
                u = boundaryCondition.conditionNames.index(self.grid.namesOrder[idx])
                coords[idx] = 0 if boundaryCondition.LeftOrRight == 'Left' else self.grid.dimSizes[idx]-1

        tempBCNodes = []
        currentDimToChange = 0
        while currentDimToChange < len(variableIndices):
            if coords[variableIndices[currentDimToChange]] < self.grid.dimSizes[variableIndices[currentDimToChange]]:
                if not coords in self.BoundaryConditionNodes:
                    tempBCNodes.append(deepcopy(coords))
                coords[variableIndices[currentDimToChange]] += 1
            else:
                coords[variableIndices[currentDimToChange]] = 0
                currentDimToChange += 1
                if currentDimToChange == len(variableIndices):
                    break
                coords[variableIndices[currentDimToChange]] += 1
                if coords[variableIndices[currentDimToChange]] < self.grid.dimSizes[variableIndices[currentDimToChange]]:
                    currentDimToChange = 0
        return tempBCNodes
        
    def applyDirichletBoundaryCondition(self, boundaryCondition):
        tempBCNodes = self.calcBoundaryConditionNodes(boundaryCondition)

        for coords in tempBCNodes:
            self.grid.setData(coords, boundaryCondition.funct(self.grid.getVariableValues(coords)))
            
        self.BoundaryConditionNodes += tempBCNodes


    def applyVanNeumanBoundaryCondition(self,boundaryCondition, relations):
        tempBCNodes = self.calcBoundaryConditionNodes(boundaryCondition)
        
        for coords in tempBCNodes:
            self.SpecialNodes.append(coords)
            self.SpecialRelations.append(relations)

        
    def doCalculation(self):
        pass
            
                